﻿#include "IDAHook.h"
#include "Logger.h"
#include "detours.h"
#include "StringManager.h"
#include "SignatureUtils.h"

// 初始化静态成员
IDAHook::QStringConstructor IDAHook::s_pOriginalQStringConstructor = nullptr;
IDAHook::QStringDestructor IDAHook::s_pOriginalQStringDestructor = nullptr;
IDAHook::QStringFromUtf8 IDAHook::s_pOriginalQStringFromUtf8 = nullptr;
IDAHook::QStringFromAsciiHelper IDAHook::s_pOriginalQStringFromAsciiHelper = nullptr;
IDAHook::QTranslate IDAHook::s_pOriginalQTranslate = nullptr;
IDAHook::QWindowSetTitle IDAHook::s_pOriginalQWindowSetTitle = nullptr;
IDAHook::QLabelSetText IDAHook::s_pOriginalQLabelSetText = nullptr;
IDAHook::QPainterDrawText IDAHook::s_pOriginalQPainterDrawText = nullptr;

// 实例化 IDAHook 模块日志
Logger IDALog(L"IDAHook.log", false, false);

bool IDAHook::Install() {
	if (!LoadIDAFunctions()) {
		IDALog.Log("[IDAHook] 错误: 加载IDA函数失败");
		return false;
	}

	// 安装所有Detours钩子
	LONG status = NO_ERROR;
	status |= DetourAttach((PVOID*)&s_pOriginalQStringFromUtf8, HookedQStringFromUtf8);
	//status |= DetourAttach((PVOID*)&s_pOriginalQStringFromAsciiHelper, HookedQStringFromAsciiHelper);
	status |= DetourAttach((PVOID*)&s_pOriginalQTranslate, HookedQTranslate);
	status |= DetourAttach((PVOID*)&s_pOriginalQWindowSetTitle, HookedQWindowSetTitle);
	status |= DetourAttach((PVOID*)&s_pOriginalQLabelSetText, HookedQLabelSetText);
	status |= DetourAttach((PVOID*)&s_pOriginalQPainterDrawText, HookedQPainterDrawText);

	if (status == NO_ERROR) {
		IDALog.Log("[IDAHook] 钩子安装成功");
		return true;
	}

	IDALog.Log("[IDAHook] 错误: 钩子安装失败，状态码: 0x%X", status);
	return false;
}

bool IDAHook::Uninstall() {
	// 检查是否有有效钩子
	if (!s_pOriginalQStringFromUtf8) return true;

	// 卸载所有Detours钩子
	LONG status = NO_ERROR;
	status |= DetourDetach((PVOID*)&s_pOriginalQStringFromUtf8, HookedQStringFromUtf8);
	//status |= DetourDetach((PVOID*)&s_pOriginalQStringFromAsciiHelper, HookedQStringFromAsciiHelper);
	status |= DetourDetach((PVOID*)&s_pOriginalQTranslate, HookedQTranslate);
	status |= DetourDetach((PVOID*)&s_pOriginalQWindowSetTitle, HookedQWindowSetTitle);
	status |= DetourDetach((PVOID*)&s_pOriginalQLabelSetText, HookedQLabelSetText);
	status |= DetourDetach((PVOID*)&s_pOriginalQPainterDrawText, HookedQPainterDrawText);

	if (status == NO_ERROR) {
		IDALog.Log("[IDAHook] 钩子卸载成功");
		return true;
	}

	IDALog.Log("[IDAHook] 错误: 钩子卸载失败，状态码: 0x%X", status);
	return false;
}

bool IDAHook::LoadIDAFunctions() {
	// 加载必要的Qt模块
	const auto LoadModule = [](const char* name) -> HMODULE {
		HMODULE module = GetModuleHandleA(name);
		if (!module) {
			IDALog.Log("[IDAHook] 错误: 未加载 %s", name);
		}
		return module;
		};

	HMODULE hQt5Core = LoadModule("Qt5Core.dll");
	HMODULE hQtGui = LoadModule("Qt5Gui.dll");
	HMODULE hQtWidgets = LoadModule("Qt5Widgets.dll");
	if (!hQt5Core || !hQtGui || !hQtWidgets) return false;

	// 获取函数地址
	s_pOriginalQStringFromUtf8 = (QStringFromUtf8)GetProcAddress(hQt5Core, "?fromUtf8@QString@QT@@SA?AV12@PEBDH@Z");
	s_pOriginalQStringFromAsciiHelper = (QStringFromAsciiHelper)GetProcAddress(hQt5Core, "?fromAscii_helper@QString@QT@@CAPEAU?$QTypedArrayData@G@2@PEBDH@Z");
	s_pOriginalQTranslate = (QTranslate)GetProcAddress(hQt5Core, "?translate@QCoreApplication@QT@@SA?AVQString@2@PEBD00H@Z");
	s_pOriginalQWindowSetTitle = (QWindowSetTitle)GetProcAddress(hQtGui, "?setTitle@QWindow@QT@@QEAAXAEBVQString@2@@Z");
	s_pOriginalQLabelSetText = (QLabelSetText)GetProcAddress(hQtWidgets, "?setText@QLabel@QT@@QEAAXAEBVQString@2@@Z");
	s_pOriginalQPainterDrawText = (QPainterDrawText)GetProcAddress(hQtGui, "?drawText@QPainter@QT@@QEAAXAEBVQRect@2@HAEBVQString@2@PEAV32@@Z");
	s_pOriginalQStringConstructor = (QStringConstructor)GetProcAddress(hQt5Core, "??0QString@QT@@QEAA@PEBVQChar@1@H@Z");
	s_pOriginalQStringDestructor = (QStringDestructor)GetProcAddress(hQt5Core, "??1QString@QT@@QEAA@XZ");

	// 验证所有函数是否加载成功
	const bool allFunctionsLoaded = s_pOriginalQStringConstructor && s_pOriginalQStringFromUtf8 &&
		s_pOriginalQStringFromAsciiHelper && s_pOriginalQTranslate && s_pOriginalQWindowSetTitle &&
		s_pOriginalQLabelSetText && s_pOriginalQPainterDrawText;


	// 使用SignatureUtils查找特定函数地址
	//std::string signatureStr = "48 89 5C 24 08 48 89 6C 24 18 48 89 74 24 20 57 48 83 EC 20 48 8B FA 48 8B D9 0F B6 42 40 F6 D0 A8 01 74 62 48 83 7A 30 00 74 5B 41 B9 41 00 00 00 45 33 C0 48 8B 15 65 EA 3F 00 48 8D 4F 38 E8 5C 2D 1F 00 0F B6 4F 40 48 8B ";
	//void* FuncAddr = SignatureUtils::DevFindSignature("ida.dll", signatureStr);

	//s_pOriginalFunc = reinterpret_cast<Func>(SignatureUtils::GetHardcodedAddress("ida.dll", 0x123A80, 0x156000));


	if (!allFunctionsLoaded) {
		IDALog.Log("[IDAHook] 错误: 部分函数地址获取失败");
		return false;
	}

	IDALog.Log("[IDAHook] 所有函数地址获取成功");
	return true;
}

void __fastcall IDAHook::HookedQStringFromUtf8(void* ret, const char* str, int size) {
	const char* finalStr = str;
	int finalSize = size;

	if (str) {
		const size_t actualSize = (size == -1) ? strlen(str) : size;
		const std::string_view originalStr(str, actualSize);

		// 只记录关键日志
		if (actualSize > 0 && actualSize < 256) {
			IDALog.Log("[QString::FromUtf8] 原始字符串: %s", originalStr.data());
			//IDALog.Log("{[QString::FromUtf8]  \"%s\", u8\"%s\"},", originalStr.data(), originalStr.data());
		}

		// 优先尝试精确匹配
		if (const char* exactMatch = StringManager::FindDefault(originalStr.data())) {
			IDALog.Log("[QString::FromUtf8] 精确匹配: %s", exactMatch);
			finalStr = exactMatch;
			finalSize = -1;
		}
		// 其次尝试模糊匹配, 长度检查：超过50字符不替换
		else if (actualSize <= 50) {
			const std::string fuzzyMatch = StringManager::FuzzyMatchDefault(originalStr.data());
			if (!fuzzyMatch.empty() && fuzzyMatch != originalStr) {
				IDALog.Log("[QString::FromUtf8] 模糊匹配: %s", fuzzyMatch);
				finalStr = fuzzyMatch.c_str();
				finalSize = -1;
			}
		}
		else {
			IDALog.Log("[QString::FromUtf8] 长度超过50字符，跳过替换");
		}
	}

	s_pOriginalQStringFromUtf8(ret, finalStr, finalSize);
}

__int64 __fastcall IDAHook::HookedQStringFromAsciiHelper(__int64 a1, int a2) {
    if (a1) {
        const char* OriginalStr = reinterpret_cast<const char*>(a1);
        IDALog.Log("[QString::FromAsciiHelper] 原始字符串: %s", OriginalStr);

        if (const char* rep = StringManager::FindDefault(OriginalStr)) {
            IDALog.Log("[QString::FromAsciiHelper] 替换: [%s] -> %s", OriginalStr, rep);
            return s_pOriginalQStringFromAsciiHelper(reinterpret_cast<__int64>(rep), a2);
        }
    }

    return s_pOriginalQStringFromAsciiHelper(a1, a2);
}

void __fastcall IDAHook::HookedQTranslate(void* ret, const char* context,
	const char* key, const char* disambiguation, int n) {

	if (key) {
		IDALog.Log("[QTranslate] 原始字符串: %s", key);
		//IDALog.Log("{[QTranslate]  \"%s\", u8\"%s\"},", key, key);
		// 精确匹配
		if (const char* exactMatch = StringManager::FindDefault(key)) {
			IDALog.Log("[QTranslate] 精确匹配: %s", exactMatch);
			s_pOriginalQTranslate(ret, context, exactMatch, disambiguation, n);
			return;
		}
	}

	// 无匹配则使用原始字符串
	s_pOriginalQTranslate(ret, context, key, disambiguation, n);
}

// 窗口标题设置钩子
void __fastcall IDAHook::HookedQWindowSetTitle(void* pThis, const void* pQString) {
	static thread_local bool isInHook = false;
	if (isInHook) {
		s_pOriginalQWindowSetTitle(pThis, pQString);
		return;
	}
	isInHook = true;

	// 提取原始字符串内容
	const std::wstring original = ExtractStringFromQString(pQString);
	void* targetString = const_cast<void*>(pQString);
	bool isModified = false;

	if (!original.empty()) {
		//IDALog.Log("[SetTitle]  {L\"%s\", L\"%s\"},", original.c_str(), original.c_str());
		IDALog.Log(L"[SetTitle] 原始窗口标题: %s", original);
		const std::wstring replaced = StringManager::FuzzyMatchDefault(original);

		// 如果找到有效替换
		if (!replaced.empty() && replaced != original) {
			if (void* tempQString = CreateTemporaryQString(replaced.c_str())) {
				targetString = tempQString;
				isModified = true;
				IDALog.Log(L"[SetTitle] 模糊匹配: %s", original);
			}
		}
	}

	// 调用原始函数
	s_pOriginalQWindowSetTitle(pThis, targetString);

	// 清理临时对象
	if (isModified) {
		ReleaseTemporaryQString(targetString);
	}

	isInHook = false;
}

// 标签文本设置钩子
void __fastcall IDAHook::HookedQLabelSetText(void* pThis, const void* pQString) {
	static thread_local bool isInHook = false;
	if (isInHook) {
		s_pOriginalQLabelSetText(pThis, pQString);
		return;
	}
	isInHook = true;

	const std::wstring original = ExtractStringFromQString(pQString);
	void* targetString = const_cast<void*>(pQString);
	bool isModified = false;

	if (!original.empty()) {
		IDALog.Log(L"[SetText] 原始标签文本: %s", original);
		//IDALog.Log("[SetText]  {L\"%s\", L\"%s\"},", original.c_str(), original.c_str());
		const std::wstring replaced = StringManager::FuzzyMatchDefault(original);

		if (!replaced.empty() && replaced != original) {
			if (void* tempQString = CreateTemporaryQString(replaced.c_str())) {
				targetString = tempQString;
				IDALog.Log(L"[SetText] 精确匹配: %s", replaced);
				isModified = true;
			}
		}
	}

	s_pOriginalQLabelSetText(pThis, targetString);

	if (isModified) {
		ReleaseTemporaryQString(targetString);
	}

	isInHook = false;
}

// 绘制文本钩子
void __fastcall IDAHook::HookedQPainterDrawText(void* pThis, void* rect, int flags, void* text, void* boundingRect) {
	static thread_local bool isInHook = false;
	if (isInHook) return;
	isInHook = true;

	const std::wstring original = ExtractStringFromQString(text);
	void* targetText = text;
	bool isModified = false;

	if (!original.empty()) {
		IDALog.Log(L"[DrawText] 原始绘制文本: %s", original);
		//IDALog.Log("[DrawText]  {L\"%s\", L\"%s\"},", original, original);
		// 优先尝试精确匹配
		if (const wchar_t* exactMatch = StringManager::FindDefault(original.c_str())) {
			if (void* tempQString = CreateTemporaryQString(exactMatch)) {
				targetText = tempQString;
				isModified = true;
			}
		}
		// 其次尝试模糊匹配
		else {
			const std::wstring fuzzyMatch = StringManager::FuzzyMatchDefault(original.c_str());
			if (!fuzzyMatch.empty() && fuzzyMatch != original) {
				if (void* tempQString = CreateTemporaryQString(fuzzyMatch.c_str())) {
					targetText = tempQString;
					isModified = true;
				}
			}
		}
	}

	// 调用原始函数
	s_pOriginalQPainterDrawText(pThis, rect, flags, targetText, boundingRect);

	// 清理临时对象
	if (isModified) {
		ReleaseTemporaryQString(targetText);
	}

	isInHook = false;
}

// 从QString提取字符串内容
std::wstring IDAHook::ExtractStringFromQString(const void* pQString) {
	if (!pQString) return L"";

	// 获取内部数据结构指针
	const uintptr_t d_ptr = *reinterpret_cast<const uintptr_t*>(pQString);
	if (!d_ptr) return L"";

	// 获取字符串长度 (使用定义的常量)
	const int length = *reinterpret_cast<const int*>(d_ptr + STRING_LENGTH_OFFSET);
	if (length <= 0 || length > 10000) return L"";

	// 获取字符串数据指针 (使用定义的常量)
	const wchar_t* pStrData = reinterpret_cast<const wchar_t*>(d_ptr + STRING_DATA_OFFSET);

	// 安全检查
	if (IsBadReadPtr(pStrData, (length + 1) * sizeof(wchar_t))) {
		return L"";
	}

	return std::wstring(pStrData, length);
}

// 修改QString内容（注意: 高风险操作）
bool IDAHook::ModifyQString(const void* pQString, const std::wstring& newContent) {
	if (!pQString || newContent.empty()) {
		return false;
	}

	// 获取d指针
	const uintptr_t d_ptr = *reinterpret_cast<const uintptr_t*>(pQString);
	if (!d_ptr) {
		return false;
	}

	// 获取当前字符串信息 (使用定义的常量)
	const int currentSize = *reinterpret_cast<const int*>(d_ptr + STRING_LENGTH_OFFSET);

	// 获取容量信息（通常位于长度字段之后）
	static constexpr uintptr_t STRING_CAPACITY_OFFSET = 0x8;  // 容量字段偏移
	const int currentCapacity = *reinterpret_cast<const int*>(d_ptr + STRING_CAPACITY_OFFSET);

	// 获取字符串数据指针 (使用定义的常量)
	wchar_t* strData = const_cast<wchar_t*>(
		reinterpret_cast<const wchar_t*>(d_ptr + STRING_DATA_OFFSET));

	// 安全检查
	if (currentSize <= 0 || currentSize > 10000 ||
		currentCapacity <= 0 || currentCapacity > 10000 ||
		currentSize > currentCapacity ||
		IsBadWritePtr(strData, (currentCapacity + 1) * sizeof(wchar_t))) {
		return false;
	}

	// 计算新内容长度（包括终止符）
	const size_t newSize = newContent.size();
	const size_t requiredCapacity = newSize + 1;  // +1 for null terminator

	// 检查缓冲区容量是否足够
	if (requiredCapacity <= static_cast<size_t>(currentCapacity)) {
		// 复制新内容
		wmemcpy(strData, newContent.c_str(), newSize);
		strData[newSize] = L'\0'; // 添加终止符

		// 更新长度字段 (使用定义的常量)
		*const_cast<int*>(reinterpret_cast<const int*>(d_ptr + STRING_LENGTH_OFFSET)) =
			static_cast<int>(newSize);

		return true;
	}

	// 新内容太大，无法放入现有缓冲区
	return false;
}

// 创建临时QString对象
void* IDAHook::CreateTemporaryQString(const wchar_t* content) {
	if (!content || !s_pOriginalQStringConstructor) {
		IDALog.Log(L"[IDAHook] 错误: 无效参数或构造函数未初始化");
		return nullptr;
	}

	// 分配内存（使用定义的QString对象大小常量）
	void* pMemory = operator new(QSTRING_OBJECT_SIZE);
	if (!pMemory) return nullptr;

	memset(pMemory, 0, QSTRING_OBJECT_SIZE);  // 清空内存

	// 调用构造函数创建对象
	s_pOriginalQStringConstructor(pMemory, content, static_cast<int>(wcslen(content)));
	return pMemory;
}

void* IDAHook::CreateTemporaryQString(const std::wstring& content) {
	return CreateTemporaryQString(content.c_str());
}

// 释放临时QString对象
void IDAHook::ReleaseTemporaryQString(void* pQString) {
	if (!pQString) return;

	// 调用析构函数
	if (s_pOriginalQStringDestructor) {
		//IDALog.Log("[IDAHook] 信息: 调用QString析构函数");
		s_pOriginalQStringDestructor(pQString);
	}

	//if (pQString) {
	//	operator delete(pQString);
	//}
}
